--[[## ft=lua ##################################################################
  This file is licensed under [WTFPL v2](https://www.wtfpl.net/).  You may read
  the full license in the file LICENSE or by following the web link.

  SET - flexible set operations in Lua

  REPRESENTATION

    A set is represented with two tables: r and s.  r is indexed by the key and
    returns an index into s.  The key is stored in s at the given index.  This
    representation is chosen for reasons of time performance.  Two tables wastes
    a bit of space, and makes insertion marginally more expensive (two O(1)
    table operations instead of one), but for that removal is also O(1).

    Insertion involves pushing the value to the end of the s table and then
    pairing that last index with the value in r.  Membership checking is as
    simple as looking up the value in r and returning true if a number is
    provided.

    Removal is a bit more involved, but for that O(1).  First membership is
    checked (O(1)).  All subsequent steps assume membership is there.  The
    value is then stubbed out of r (O(1)).  The last value is popped off of the
    table s (O(1)).  If that value doesn't match the erased key, r is indexed
    by the popped value and given the value of where s contained the removed
    value (O(1)), followed by s[i] being given the value that was popped (O(1)).

    Thus, with this representation, at a slight cost in memory, sets have O(1)
    insertion, O(1) membership tests, and O(1) removal (albeit at wildly
    different K factors).

    Note that the set exposes its members (in arbitrary order) as a read-only,
    array-flavoured table.  This permits it to be used directly in most table
    reading operations.  Changing it with rawset is possible but will almost
    certainly break the set operations.
--############################################################################]]

--[[###################################
    Constructor.  Accepts any number of arbitrary arguments as the starting values
    of the set.
--###################################]]
local function new(_, ...)
    local r = {}
    local s = {}
    local protected = true
    local instantiated = false

    local function is_set_impl(s)
        return getmetatable(s) and getmetatable(s).__type__ == 'set'
    end

    setmetatable(s, {
        -- API operations
        __index = {
            -- single value insertion
            insert = function(ss, v)
                assert(is_set_impl(ss), 'self is not a set')
                assert(v ~= 'nil', 'cannot have nil as member')
                if not r[v] then
                    protected = false
                    table.insert(ss, v)
                    protected = true
                    r[v] = #ss
                end
                return ss
            end,

            -- single value removal
            remove = function(ss, v)
                assert(is_set_impl(ss), 'self is not a set')
                local i = r[v]
                if i then
                    r[v] = nil
                    protected = false
                    local t = table.remove(ss)
                    protected = true
                    if t ~= v then
                        r[t] = i
                        ss[i] = t
                    end
                end
                return ss
            end,

            -- multi-value insertion
            batch_insert = function(ss, ...)
                assert(is_set_impl(ss), 'self is not a set')
                for _,v in pairs {...} do
                    ss:insert(v)
                end
                return ss
            end,

            -- multi-value removal
            batch_remove = function(ss, ...)
                assert(is_set_impl(ss), 'self is not a set')
                for _,v in pairs {...} do
                    ss:remove(v)
                end
                return ss
            end,

            -- membership test
            has = function(ss, e)
                assert(is_set_impl(ss), 'self is not a set')
                return r[e] ~= nil
            end,
        },

        -- set manipulation operators

        -- union
        __add = function(s1, s2)
            assert(is_set_impl(s1), 'first argument is not a set')
            assert(is_set_impl(s2), 'second argument is not a set')
            local r = new()
            for _,v in ipairs(s1) do r:insert(v) end
            for _,v in ipairs(s2) do r:insert(v) end
            return r
        end,

        -- subtraction
        __sub = function(s1, s2)
            assert(is_set_impl(s1), 'first argument is not a set')
            assert(is_set_impl(s2), 'second argument is not a set')
            local r = new()
            for _,v in ipairs(s1) do r:insert(v) end
            for _,v in ipairs(s2) do r:remove(v) end
            return r
        end,

        -- intersection
        __mul = function(s1, s2)
            assert(is_set_impl(s1), 'first argument is not a set')
            assert(is_set_impl(s2), 'second argument is not a set')
            local r = new()
            for _,v in ipairs(s1) do
                if s2:has(v) then
                    r:insert(v)
                end
            end
            return r
        end,

        -- equality
        __eq = function(s1, s2)
            assert(is_set_impl(s1), 'first argument is not a set')
            assert(is_set_impl(s2), 'second argument is not a set')
            if #s1 ~= #s2 then return false
            else
                for _,v in ipairs(s1) do
                    if not s2:has(v) then return false end
                end
            end
            return true
        end,

        -- proper subset
        __lt = function(s1, s2)
            assert(is_set_impl(s1), 'first argument is not a set')
            assert(is_set_impl(s2), 'second argument is not a set')
            if s1 == s2 then return false
            else
                for _,v in ipairs(s1) do
                    if not s2:has(v) then return false end
                end
            end
            return true
        end,

        -- subset
        __le = function(s1, s2)
            assert(is_set_impl(s1), 'first argument is not a set')
            assert(is_set_impl(s2), 'second argument is not a set')
            return (s1 == s2) or (s1 < s2)
        end,

        -- write protection
        __newindex = function(ss, k, v)
            assert(is_set_impl(ss), 'first argument is not a set')
            if protected then
                error('attempting to write directly to set', 2)
            else
                rawset(ss, k, v)
            end
        end,

        -- metatable type tag
        __type__ = 'set'
    })
    s:batch_insert(...)
    instantiated = true
    return s
end

return setmetatable({}, { __call = new })
